from .vit import VisionTransformer as vvit
from .vit_moe import VisionTransformer as vmoe
from .vit_multi import VisionTransformer as vmulti
from .vit_ssiat import VisionTransformer as vssiat
# from .vit_ease import VisionTransformer as vease
# from .vit_sdm import VisionTransformer as vsdm
# from .vit_baseSingle import VisionTransformer as vbase_single
# from .vit_baseMulti import VisionTransformer as vbase_multi
# from .vit_multi_mnc import VisionTransformer as vmmnc
# from .vit_multi_wbb import VisionTransformer as vmulti_wbb
# from .vit_multi_SideNet import VisionTransformer as vmulti_sn
# from .vit_multi_downSideNet import VisionTransformer as vmulti_dsn
# from .vit_multi_UniPT import VisionTransformer as vmulti_unipt
# from .vit_multi_blocks import VisionTransformer as vmulti_blocks


from typing import Union, cast, Type

VIT_TYYPE = Union[
    vvit,
    vmoe,
    vmulti,
    vssiat,
    # vease,
    # vsdm,
    # vbase_single,
    # vbase_multi,
    # vmmnc,
    # vmulti_wbb,
    # vmulti_sn,
    # vmulti_dsn,
    # vmulti_unipt,
    # vmulti_blocks,
]

__all__ = [
    "vvit",
    "vmoe",
    "vmulti",
    "vssiat",
    # "vease",
    # "vsdm",
    # "vbase_single",
    # "vbase_multi",
    # "vmmnc",
    # "vmulti_wbb",
    # "vmulti_sn",
    # "vmulti_dsn",
    # "vmulti_unipt",
    # "vmulti_blocks",
]

import timm
from utils.toolkit import NamespaceDict


def get_tuning_config(args, extra_params=None):
    # FIXME: too many config mixed in args
    config = {
        "embed_dim": 768,
        "peft_name": args.get("peft_name", None),
        "ffn_option": args.get("ffn_option", "none"),
        "ffn_rank": args.get("ffn_rank", None),
        "ffn_adapter_scalar": args.get("ffn_adapter_scalar", 1.0),
        "ffn_adapter_layernorm_option": args.get(
            "ffn_adapter_layernorm_option", "none"
        ),
        "mnc_norm": args.get("mnc_norm", False),
        "moe_topk": args.get("moe_topk", None),
        "cam_visual": args["cam_visual"],
        "_device": args["device"][0],
    }
    config.update({"sidenet_cfg": args.get("sidenet_cfg", None)})
    if extra_params:
        config.update(extra_params)
    return NamespaceDict(config)


def get_backbone(
    args,
) -> VIT_TYYPE:
    """
    返回一个VisionTransformer的子类实例
    返回类型取决于args中指定的backbone_type
    所有可能的返回类型都是vvit的子类
    """
    name = args["backbone_type"]

    config = get_tuning_config(args)
    model = timm.create_model(
        name,
        num_classes=0,
        global_pool=False,
        drop_path_rate=0.0,
        config=config,
    )
    model = cast(VIT_TYYPE, model)
    model.out_dim = 768
    return model.eval()
